package com.example.footclassfiydome.service;

import ai.djl.Model;
import ai.djl.basicdataset.cv.classification.ImageFolder;
import ai.djl.metric.Metrics;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.ndarray.types.Shape;
import ai.djl.training.DefaultTrainingConfig;
import ai.djl.training.EasyTrain;
import ai.djl.training.Trainer;
import ai.djl.training.dataset.RandomAccessDataset;
import ai.djl.training.evaluator.Accuracy;
import ai.djl.training.listener.TrainingListener;
import ai.djl.training.loss.Loss;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import org.springframework.stereotype.Service;

import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 * Created by lilijava
 * on 2024/12/25
 * <p>
 * on  foot-classfiy-dome
 * @author DELL
 */
@Service
public class TrainService {
    public void train (String datasetPath,String modelPath) throws IOException, TranslateException {
        //准备数据集
        ImageFolder dataset = initDataset(Paths.get(datasetPath));
        //2分割数据集
        RandomAccessDataset[] split = dataset.randomSplit(8, 2);
        //训练集
        RandomAccessDataset trainDataset = split[0];
        //验证集
        RandomAccessDataset validationDataset = split[1];
        //定义模型
        try(Model model = Models.getModel()){
            //准备训练配置，获取训练器
            DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
                    .addEvaluator(new Accuracy())
                    .addTrainingListeners(TrainingListener.Defaults.logging());
            //获取训练器
            Trainer trainer = model.newTrainer(config);
            trainer.setMetrics(new Metrics());
            //初始化（初始用一个图片，先启动）
            trainer.initialize(new Shape(1,3, 100, 100));
            //训练
            EasyTrain.fit(trainer, 2, trainDataset, validationDataset);
            //保存模型
            model.save(Paths.get(modelPath), "footweaver");

            //保存 synset.txt 文件
            List<String> synset = dataset.getSynset();
            // synset.txt
            Path modelDir = Paths.get(modelPath);
            Path synsetPath = modelDir.resolve("synset.txt");
            try(BufferedWriter writer = Files.newBufferedWriter(synsetPath)){
                writer.write(String.join("\n",synset));
            }
        }
    }
//初始化数据集
        public static ImageFolder initDataset(Path path) throws IOException {
            ImageFolder dataset = ImageFolder.builder()
                    .setRepositoryPath(path)
                    .setSampling(128, true)
                    .optMaxDepth(10)
                    .addTransform(new Resize(100, 100))
                    .addTransform(new ToTensor())
                    .build();
            dataset.prepare(new ProgressBar());
            return dataset;
        }
    }
